#!/usr/bin/env python3
"""
Debug CUDA device-side assertion error
Run with detailed error reporting
"""
import os
import sys

# Enable synchronous CUDA execution for accurate error location
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

# Enable device-side assertions for detailed error messages
os.environ['TORCH_USE_CUDA_DSA'] = '1'

print("=" * 80)
print("CUDA Debug Mode Enabled")
print("=" * 80)
print("CUDA_LAUNCH_BLOCKING=1")
print("TORCH_USE_CUDA_DSA=1")
print("=" * 80)

import torch
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA version: {torch.version.cuda}")
    print(f"Device name: {torch.cuda.get_device_name(0)}")
    print(f"Device count: {torch.cuda.device_count()}")
print("=" * 80)

# Try to initialize the model
print("\nAttempting to initialize ImprovedROIRefinerModel...")
print("=" * 80)

try:
    from src.models.refiner_improved import ImprovedROIRefinerModel
    from src.config import IMPROVED_MODEL_CONFIG

    print("Creating model instance...")
    # ImprovedROIRefinerModel reads num_classes from data.yaml automatically
    # No need to pass it as parameter
    model = ImprovedROIRefinerModel(
        device='cuda',
        unfreeze_layers=2,
        config=IMPROVED_MODEL_CONFIG
    )
    print("✓ Model created successfully!")

except Exception as e:
    print(f"\n{'='*80}")
    print("ERROR CAUGHT:")
    print(f"{'='*80}")
    print(f"Type: {type(e).__name__}")
    print(f"Message: {str(e)}")
    print(f"{'='*80}")

    import traceback
    print("\nFull traceback:")
    traceback.print_exc()

    print(f"\n{'='*80}")
    print("DEBUGGING HINTS:")
    print(f"{'='*80}")

    if "assert" in str(e).lower():
        print("✗ Device-side assertion triggered")
        print("  Possible causes:")
        print("  1. Index out of bounds in tensor operations")
        print("  2. Invalid memory access")
        print("  3. Tensor shape mismatch")
        print("  4. NaN or Inf values")
        print("\nNext steps:")
        print("  - Check OverLock model initialization")
        print("  - Verify tensor shapes in cross_roi_attention")
        print("  - Check for any hardcoded indices")

    sys.exit(1)

print("\n" + "=" * 80)
print("SUCCESS: No errors detected!")
print("=" * 80)
